{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}

module Control.Fold1 (Fold1, fold1) where

import Control.Applicative (liftA2)
import Data.Biapplicative (biliftA2)
import Data.Bifunctor (bimap, first, second)
import Data.Kind (Type)
import Data.Profunctor (Choice (left', right'), Profunctor (dimap, lmap, rmap))
import Data.Semigroup.Foldable (Foldable1)

data Fold1 (a :: Type) (b :: Type) where
  Fold1 :: (x -> x -> x) -> (a -> x) -> (x -> b) -> Fold1 a b

instance Functor (Fold1 a) where
  {-# INLINEABLE fmap #-}
  fmap f (Fold1 step inj proj) = Fold1 step inj (fmap f proj)
  {-# INLINEABLE (<$) #-}
  x <$ _ = pure x

instance Applicative (Fold1 a) where
  {-# INLINEABLE pure #-}
  pure x = Fold1 (\_ _ -> ()) (const ()) (const x)
  {-# INLINEABLE (<*>) #-}
  Fold1 stepF injF projF <*> Fold1 stepX injX projX =
    Fold1
      (biliftA2 stepF stepX)
      (\x -> (injF x, injX x))
      (uncurry ($) . bimap projF projX)
  {-# INLINEABLE liftA2 #-}
  liftA2 f (Fold1 stepX injX projX) (Fold1 stepY injY projY) =
    Fold1
      (biliftA2 stepX stepY)
      (\x -> (injX x, injY x))
      (uncurry f . bimap projX projY)
  {-# INLINEABLE (*>) #-}
  _ *> x = x
  {-# INLINEABLE (<*) #-}
  x <* _ = x

instance Profunctor Fold1 where
  {-# INLINEABLE dimap #-}
  dimap f g (Fold1 step inj proj) = Fold1 step (inj . f) (g . proj)
  {-# INLINEABLE lmap #-}
  lmap f (Fold1 step inj proj) = Fold1 step (inj . f) proj
  {-# INLINEABLE rmap #-}
  rmap g (Fold1 step inj proj) = Fold1 step inj (g . proj)

instance Choice Fold1 where
  {-# INLINEABLE left' #-}
  left' (Fold1 step inj proj) =
    Fold1
      ( \acc x -> case acc of
          Left y -> case x of
            Left x' -> Left (step y x')
            _ -> x
          _ -> acc
      )
      (first inj)
      (first proj)
  {-# INLINEABLE right' #-}
  right' (Fold1 step inj proj) =
    Fold1
      (liftA2 step)
      (second inj)
      (second proj)

instance (Semigroup b) => Semigroup (Fold1 a b) where
  {-# INLINEABLE (<>) #-}
  (<>) = liftA2 (<>)

instance (Monoid b) => Monoid (Fold1 a b) where
  {-# INLINEABLE mempty #-}
  mempty = pure mempty

fold1 :: (Functor f, Foldable1 f) => Fold1 a b -> f a -> b
fold1 (Fold1 step inj proj) = proj . foldl1 step . fmap inj
